/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/

#include "environment.hxx"

int main(int argc, char **argv) {

  linal::initialize( argc, argv );

  bool help;
  int nthreads, datatype, m, n, n_rhs, b_mn;
  double tol;

  linal::get_option_bool  ( "-help", "list options",
                            __FILE__,LINAL_OPTION_MODE_RUN, help);

  linal::get_option_int   ( "-datatype", LINAL_OPTION_TEXT_DATATYPE,
                            __FILE__,TEST_DATATYPE, datatype, help);

  linal::get_option_int   ( "-M", "size of matrix",
                            __FILE__,N, m, help);
  linal::get_option_int   ( "-N", "size of matrix",
                            __FILE__,N, n, help);
  linal::get_option_int   ( "-Nrhs", "# of rhs",
                            __FILE__,1, n_rhs, help);

  linal::get_option_int   ( "-B_MN", "block size of matrix",
                            __FILE__,B_MN, b_mn, help);

  linal::get_option_int   ( "-nthreads", "# of threads",
                            __FILE__,1, nthreads, help);

  linal::get_option_double( "-tol", "error tolerence",
                            __FILE__,LINAL_ERROR_TOL, tol, help);

  if (help == LINAL_OPTION_MODE_DRY) {
    linal::finalize();
    return 0;
  }

  omp_set_num_threads(nthreads);

  linal::Flat_ A[2], T[2], W[2], X[2], B, Y, norm;
  linal::Hier_ HA, HT, HW, HX, HY;

  // m > n
  for (int i=0;i<2;++i) {
    A[i].create( datatype, m, n );
    T[i].create( datatype, n, n );
    W[i].create( datatype, n, n );
    X[i].create( datatype, n, n_rhs );
  }
  B.create( datatype, m, n_rhs );
  Y.create( datatype, m, n_rhs );
  norm.create( linal::get_datatype_real(datatype), 1, 1 );

  HA.create( A[0], b_mn, b_mn );
  HT.create( T[0], b_mn, b_mn );
  HW.create( W[0], b_mn, b_mn );
  HX.create( X[0], b_mn, b_mn );
  HY.create( Y, b_mn, b_mn );

  FLA_Random_matrix( ~A[0] );
  FLA_Random_matrix( ~B );
  FLA_Copy( ~A[0], ~A[1] );

  if (linal::get_verbose()) {
    FLA_Obj_show( "= A[0] =", ~A[0], "% 4.3f", "===" );
    FLA_Obj_show( "= A[1] =", ~A[1], "% 4.3f", "===" );
  }

#pragma omp parallel
  {
#pragma omp single nowait
      linal::qr( HA, HT );
  }

  FLA_Copy( ~B, ~Y );
  linal::apply_q( FLA_LEFT, FLA_TRANSPOSE, FLA_FORWARD, FLA_COLUMNWISE,
                  HA, HT, HW, HY );

  FLA_Obj AT, AB, YT, YB;
  FLA_Part_2x1( ~A[0], &AT,
                       &AB, FLA_Obj_width( ~A[0] ), FLA_TOP );
  FLA_Part_2x1( ~Y, &YT,
                    &YB, FLA_Obj_width( ~A[0] ), FLA_TOP );

  FLA_Trsm( FLA_LEFT, FLA_UPPER_TRIANGULAR, FLA_NO_TRANSPOSE,
            FLA_NONUNIT_DIAG, FLA_ONE, AT, YT );

  FLA_Copy( YT, ~X[0] );

  FLA_QR_UT( ~A[1], ~T[1] );
  FLA_QR_UT_solve( ~A[1], ~T[1], ~B, ~X[1] );

  if (linal::get_verbose()) {
    FLA_Obj_show( "= A[0] =", ~A[0], "% 4.3f", "===" );
    FLA_Obj_show( "= A[1] =", ~A[1], "% 4.3f", "===" );

    FLA_Obj_show( "= T[0] =", ~T[0], "% 4.3f", "===" );
    FLA_Obj_show( "= T[1] =", ~T[1], "% 4.3f", "===" );

    FLA_Obj_show( "= X[0] =", ~X[0], "% 4.3f", "===" );
    FLA_Obj_show( "= X[1] =", ~X[1], "% 4.3f", "===" );
  }

  float norm_f = 0;
  double norm_d = 0;

  FLA_Axpy( FLA_MINUS_ONE, ~A[0], ~A[1] );
  FLA_Norm1( ~A[1], ~norm );

  if (linal::get_verbose())
    FLA_Obj_show( "=Diff A[1]=", ~A[1], "% 4.3f", "===" );

  norm_f += *((float *)FLA_Obj_buffer_at_view( ~norm ));
  norm_d += *((double*)FLA_Obj_buffer_at_view( ~norm ));

  FLA_Axpy( FLA_MINUS_ONE, ~T[0], ~T[1] );
  FLA_Norm1( ~T[1], ~norm );

  if (linal::get_verbose())
    FLA_Obj_show( "=Diff T[1]=", ~T[1], "% 4.3f", "===" );

  norm_f += *((float *)FLA_Obj_buffer_at_view( ~norm ));
  norm_d += *((double*)FLA_Obj_buffer_at_view( ~norm ));

  FLA_Axpy( FLA_MINUS_ONE, ~X[0], ~X[1] );
  FLA_Norm1( ~X[1], ~norm );

  if (linal::get_verbose())
    FLA_Obj_show( "=Diff X[1]=", ~X[1], "% 4.3f", "===" );

  norm_f += *((float *)FLA_Obj_buffer_at_view( ~norm ));
  norm_d += *((double*)FLA_Obj_buffer_at_view( ~norm ));

  printf("- TEST::");
  for (int i=0;i<argc;++i)
    printf(" %s ", argv[i]);
  printf("\n");

  int rval;
  switch (linal::get_datatype_real(datatype)) {
  case LINAL_SINGLE_REAL:
    printf("SINGLE PRECISION\n");
    if (norm_f < LINAL_ERROR_TOL) {
      printf("PASS::Norm :: %10.9f \n", norm_f);   rval = 0;
    } else {
      printf("FAIL::Norm :: %10.9f \n", norm_f);   rval = -1;
    }
    break;
  case LINAL_DOUBLE_REAL:
    printf("DOUBLE PRECISION\n");
    if (norm_d < LINAL_ERROR_TOL) {
      printf("PASS::Norm :: %E \n", norm_d);   rval = 0;
    } else {
      printf("FAIL::Norm :: %E \n", norm_d);   rval = -1;
    }
    break;
  }

  HA.free();  HT.free();
  HW.free();  HX.free();
  HY.free();

  B.free();   Y.free();
  norm.free();

  for (int i=0;i<2;++i) {
    A[i].free(); T[i].free();
    W[i].free(); X[i].free();
  }
    
  linal::finalize();

  return rval;
}
